import sys
import os
import json
import redis
import time
import shutil
from datetime import datetime
from PySide6.QtWidgets import (QApplication, QMainWindow, QTabWidget, QWidget, QVBoxLayout, 
                             QHBoxLayout, QLabel, QPushButton, QListWidget, QListWidgetItem, 
                             QFileDialog, QMessageBox, QComboBox, QLineEdit, QTextEdit, 
                             QDialog, QFormLayout, QSpinBox, QDoubleSpinBox, QGridLayout,
                             QTableWidget, QTableWidgetItem, QHeaderView, QInputDialog,
                             QToolBar, QStatusBar, QFrame, QSplitter, QGroupBox)
from PySide6.QtCore import Qt, QThread, Signal, QTimer, QDateTime, QPoint
from PySide6.QtGui import QPixmap, QImage, QPainter, QPen, QColor, QIcon, QCursor, QAction

# 常量定义
USER_ROLES = ["管理员", "标注员", "检查员"]
DETECTION_TYPES = ["ocr", "分类", "rect", "四边形", "多边形"]
FILE_TYPES = ["图片", "视频"]
JOB_STATUS = ["未分配", "已分配", "已经浏览", "已标注", "检查完成", "检查错误"]
COMMENT_TYPES = ["正确", "无标签", "目标不完整", "标签不正确", "不标准"]

# 用户管理器 - 处理用户添加、删除等功能
class UserManager:
    def __init__(self, redis_host='localhost', redis_port=6379, redis_password=None):
        self.redis_client = redis.Redis(
            host=redis_host, 
            port=redis_port, 
            password=redis_password,
            decode_responses=True
        )
        self.users_key = "users"
    
    def get_all_users(self):
        """获取所有用户"""
        return self.redis_client.hgetall(self.users_key)
    
    def add_user(self, username, password, role):
        """添加用户"""
        if self.redis_client.hget(self.users_key, username):
            return False, "用户已存在"
        
        user_data = {
            "password": password,
            "role": role,
            "created_at": datetime.now().isoformat()
        }
        
        self.redis_client.hset(self.users_key, username, json.dumps(user_data))
        return True, "用户添加成功"
    
    def delete_user(self, username):
        """删除用户"""
        if not self.redis_client.hget(self.users_key, username):
            return False, "用户不存在"
        
        self.redis_client.hdel(self.users_key, username)
        return True, "用户删除成功"
    
    def authenticate(self, username, password):
        """用户认证"""
        user_data = self.redis_client.hget(self.users_key, username)
        if not user_data:
            return False, "用户不存在"
        
        user_data = json.loads(user_data)
        if user_data["password"] != password:
            return False, "密码错误"
        
        return True, user_data["role"]

# 数据管理器 - 处理与Redis和本地文件的交互

# 登录对话框（改进版，使用UserManager认证）
class LoginDialog(QDialog):
    def __init__(self, user_manager, parent=None):
        super().__init__(parent)
        self.user_manager = user_manager
        self.setWindowTitle("用户登录")
        self.resize(300, 200)
        
        layout = QVBoxLayout(self)
        
        # 用户名和密码输入
        form_layout = QFormLayout()
        self.username_edit = QLineEdit()
        self.password_edit = QLineEdit()
        self.password_edit.setEchoMode(QLineEdit.Password)
        
        form_layout.addRow("用户名:", self.username_edit)
        form_layout.addRow("密码:", self.password_edit)
        
        layout.addLayout(form_layout)
        
        # 按钮
        button_layout = QHBoxLayout()
        self.login_button = QPushButton("登录")
        self.cancel_button = QPushButton("取消")
        
        button_layout.addWidget(self.login_button)
        button_layout.addWidget(self.cancel_button)
        
        layout.addLayout(button_layout)
        
        # 连接信号和槽
        self.login_button.clicked.connect(self.validate_login)
        self.cancel_button.clicked.connect(self.reject)
    
    def validate_login(self):
        username = self.username_edit.text()
        password = self.password_edit.text()
        
        if not username or not password:
            QMessageBox.warning(self, "警告", "用户名和密码不能为空")
            return
        
        success, result = self.user_manager.authenticate(username, password)
        if success:
            self.user_role = result
            self.accept()
        else:
            QMessageBox.warning(self, "登录失败", result)

# 用户管理对话框 - 管理员添加、删除用户
class UserManagerDialog(QDialog):
    def __init__(self, user_manager, parent=None):
        super().__init__(parent)
        self.user_manager = user_manager
        self.setWindowTitle("用户管理")
        self.resize(500, 400)
        
        layout = QVBoxLayout(self)
        
        # 用户列表
        self.user_table = QTableWidget()
        self.user_table.setColumnCount(3)
        self.user_table.setHorizontalHeaderLabels(["用户名", "角色", "操作"])
        self.user_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.Stretch)
        self.user_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeToContents)
        self.user_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeToContents)
        layout.addWidget(self.user_table)
        
        # 添加用户按钮
        add_button = QPushButton("添加用户")
        add_button.clicked.connect(self.add_user)
        layout.addWidget(add_button)
        
        # 刷新用户列表
        self.refresh_users()
    
    def refresh_users(self):
        self.user_table.setRowCount(0)
        users = self.user_manager.get_all_users()
        
        for row, (username, data) in enumerate(users.items()):
            data = json.loads(data)
            self.user_table.insertRow(row)
            
            # 用户名
            self.user_table.setItem(row, 0, QTableWidgetItem(username))
            
            # 角色
            self.user_table.setItem(row, 1, QTableWidgetItem(data["role"]))
            
            # 操作按钮
            delete_button = QPushButton("删除")
            delete_button.clicked.connect(lambda checked, u=username: self.delete_user(u))
            self.user_table.setCellWidget(row, 2, delete_button)
    
    def add_user(self):
        dialog = AddUserDialog(self)
        if dialog.exec():
            username, password, role = dialog.get_user_data()
            success, message = self.user_manager.add_user(username, password, role)
            QMessageBox.information(self, "结果", message)
            self.refresh_users()
    
    def delete_user(self, username):
        reply = QMessageBox.question(
            self, "确认", f"确定要删除用户 '{username}' 吗？",
            QMessageBox.Yes | QMessageBox.No
        )
        
        if reply == QMessageBox.Yes:
            success, message = self.user_manager.delete_user(username)
            QMessageBox.information(self, "结果", message)
            self.refresh_users()

# 添加用户对话框
class AddUserDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("添加用户")
        self.resize(300, 200)
        
        layout = QVBoxLayout(self)
        
        # 表单布局
        form_layout = QFormLayout()
        
        self.username_edit = QLineEdit()
        self.password_edit = QLineEdit()
        self.password_edit.setEchoMode(QLineEdit.Password)
        self.role_combo = QComboBox()
        self.role_combo.addItems(USER_ROLES[1:])  # 管理员只能由初始用户创建
        
        form_layout.addRow("用户名:", self.username_edit)
        form_layout.addRow("密码:", self.password_edit)
        form_layout.addRow("角色:", self.role_combo)
        
        layout.addLayout(form_layout)
        
        # 按钮
        button_layout = QHBoxLayout()
        self.ok_button = QPushButton("确定")
        self.cancel_button = QPushButton("取消")
        
        self.ok_button.clicked.connect(self.accept)
        self.cancel_button.clicked.connect(self.reject)
        
        button_layout.addWidget(self.ok_button)
        button_layout.addWidget(self.cancel_button)
        
        layout.addLayout(button_layout)
    
    def get_user_data(self):
        return (
            self.username_edit.text(),
            self.password_edit.text(),
            self.role_combo.currentText()
        )

# 模型管理对话框
class ModelManagerDialog(QDialog):
    def __init__(self, data_manager, parent=None):
        super().__init__(parent)
        self.data_manager = data_manager
        self.setWindowTitle("模型管理")
        self.resize(500, 400)
        
        layout = QVBoxLayout(self)
        
        # 模型列表
        self.model_table = QTableWidget()
        self.model_table.setColumnCount(4)
        self.model_table.setHorizontalHeaderLabels(["模型名称", "标签", "检测类型", "操作"])
        self.model_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.Stretch)
        self.model_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.Stretch)
        self.model_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeToContents)
        self.model_table.horizontalHeader().setSectionResizeMode(3, QHeaderView.ResizeToContents)
        layout.addWidget(self.model_table)
        
        # 添加模型按钮
        add_button = QPushButton("添加模型")
        add_button.clicked.connect(self.add_model)
        layout.addWidget(add_button)
        
        # 刷新模型列表
        self.refresh_models()
    
    def refresh_models(self):
        self.model_table.setRowCount(0)
        models = self.data_manager.get_all_models()
        
        for row, (model_name, data) in enumerate(models.items()):
            data = json.loads(data)
            self.model_table.insertRow(row)
            
            # 模型名称
            self.model_table.setItem(row, 0, QTableWidgetItem(model_name))
            
            # 标签
            labels_str = ", ".join(data["labels"])
            self.model_table.setItem(row, 1, QTableWidgetItem(labels_str))
            
            # 检测类型
            self.model_table.setItem(row, 2, QTableWidgetItem(data["det_type"]))
            
            # 操作按钮
            delete_button = QPushButton("删除")
            delete_button.clicked.connect(lambda checked, m=model_name: self.delete_model(m))
            self.model_table.setCellWidget(row, 3, delete_button)
    
    def add_model(self):
        dialog = AddModelDialog(self)
        if dialog.exec():
            model_name, labels, det_type = dialog.get_model_data()
            self.data_manager.add_model(model_name, labels, det_type)
            self.refresh_models()
    
    def delete_model(self, model_name):
        reply = QMessageBox.question(
            self, "确认", f"确定要删除模型 '{model_name}' 吗？",
            QMessageBox.Yes | QMessageBox.No
        )
        
        if reply == QMessageBox.Yes:
            self.data_manager.redis_client.hdel(self.data_manager.models_key, model_name)
            self.refresh_models()

# 添加模型对话框
class AddModelDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("添加模型")
        self.resize(300, 200)
        
        layout = QVBoxLayout(self)
        
        # 表单布局
        form_layout = QFormLayout()
        
        self.model_name_edit = QLineEdit()
        self.labels_edit = QLineEdit()
        self.labels_edit.setPlaceholderText("多个标签用逗号分隔")
        self.det_type_combo = QComboBox()
        self.det_type_combo.addItems(DETECTION_TYPES)
        
        form_layout.addRow("模型名称:", self.model_name_edit)
        form_layout.addRow("标签:", self.labels_edit)
        form_layout.addRow("检测类型:", self.det_type_combo)
        
        layout.addLayout(form_layout)
        
        # 按钮
        button_layout = QHBoxLayout()
        self.ok_button = QPushButton("确定")
        self.cancel_button = QPushButton("取消")
        
        self.ok_button.clicked.connect(self.accept)
        self.cancel_button.clicked.connect(self.reject)
        
        button_layout.addWidget(self.ok_button)
        button_layout.addWidget(self.cancel_button)
        
        layout.addLayout(button_layout)
    
    def get_model_data(self):
        model_name = self.model_name_edit.text()
        labels = self.labels_edit.text().split(",") if self.labels_edit.text() else []
        det_type = self.det_type_combo.currentText()
        return model_name, labels, det_type

# 创建工作对话框
class CreateJobDialog(QDialog):
    def __init__(self, data_manager, parent=None):
        super().__init__(parent)
        self.data_manager = data_manager
        self.setWindowTitle("创建工作")
        self.resize(300, 300)
        
        layout = QVBoxLayout(self)
        
        # 表单布局
        form_layout = QFormLayout()
        
        self.name_edit = QLineEdit()
        self.file_type_combo = QComboBox()
        self.file_type_combo.addItems(FILE_TYPES)
        self.model_combo = QComboBox()
        self.base_url_edit = QLineEdit()
        self.base_url_button = QPushButton("选择目录")
        self.base_url_button.clicked.connect(self.select_base_url)
        
        form_layout.addRow("工作名称:", self.name_edit)
        form_layout.addRow("文件类型:", self.file_type_combo)
        form_layout.addRow("模型:", self.model_combo)
        form_layout.addRow("文件目录:", self.base_url_edit)
        form_layout.addRow("", self.base_url_button)
        
        layout.addLayout(form_layout)
        
        # 按钮
        button_layout = QHBoxLayout()
        self.ok_button = QPushButton("确定")
        self.cancel_button = QPushButton("取消")
        
        self.ok_button.clicked.connect(self.accept)
        self.cancel_button.clicked.connect(self.reject)
        
        button_layout.addWidget(self.ok_button)
        button_layout.addWidget(self.cancel_button)
        
        layout.addLayout(button_layout)
        
        # 刷新模型列表
        self.refresh_models()
    
    def select_base_url(self):
        directory = QFileDialog.getExistingDirectory(self, "选择文件目录")
        if directory:
            self.base_url_edit.setText(directory)
    
    def refresh_models(self):
        models = self.data_manager.get_all_models()
        self.model_combo.clear()
        for model_name in models.keys():
            self.model_combo.addItem(model_name)
    
    def get_job_data(self):
        name = self.name_edit.text()
        file_type = self.file_type_combo.currentText()
        model_name = self.model_combo.currentText()
        base_url = self.base_url_edit.text()
        
        job_data = {
            "name": name,
            "file_type": file_type,
            "model_name": model_name,
            "updatetime": datetime.now().isoformat(),
            "base_url": base_url
        }
        return job_data

# 分配文件对话框
class AssignFilesDialog(QDialog):
    def __init__(self, data_manager, job, parent=None):
        super().__init__(parent)
        self.data_manager = data_manager
        self.job = job
        self.setWindowTitle("分配文件")
        self.resize(500, 400)
        
        layout = QVBoxLayout(self)
        
        # 用户列表
        self.user_table = QTableWidget()
        self.user_table.setColumnCount(2)
        self.user_table.setHorizontalHeaderLabels(["用户名", "分配数量"])
        self.user_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.Stretch)
        self.user_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeToContents)
        layout.addWidget(self.user_table)
        
        # 刷新用户列表
        self.refresh_users()
        
        # 按钮
        button_layout = QHBoxLayout()
        self.ok_button = QPushButton("确定")
        self.cancel_button = QPushButton("取消")
        
        self.ok_button.clicked.connect(self.accept)
        self.cancel_button.clicked.connect(self.reject)
        
        button_layout.addWidget(self.ok_button)
        button_layout.addWidget(self.cancel_button)
        
        layout.addLayout(button_layout)
    
    def refresh_users(self):
        users = self.data_manager.user_manager.get_all_users()
        annotators = [username for username, data in users.items() if json.loads(data)["role"] == "标注员"]
        self.user_table.setRowCount(len(annotators))
        
        for row, username in enumerate(annotators):
            self.user_table.setItem(row, 0, QTableWidgetItem(username))
            spin_box = QSpinBox()
            spin_box.setRange(0, len(self.job["files"]))
            self.user_table.setCellWidget(row, 1, spin_box)
    
    def get_assigned_files(self):
        user_files = {}
        for row in range(self.user_table.rowCount()):
            username = self.user_table.item(row, 0).text()
            count = self.user_table.cellWidget(row, 1).value()
            if count > 0:
                user_files[username] = []
                files = self.data_manager.get_job_files(self.job["id"], status="未分配")
                for i in range(count):
                    if i < len(files):
                        user_files[username].append(files[i]["filename"])
        return user_files

# 管理员界面 - 增强版，添加用户管理功能
class AdminWidget(QWidget):
    def __init__(self, data_manager, parent=None):
        super().__init__(parent)
        self.data_manager = data_manager
        self.user_manager = data_manager.user_manager
        self.init_ui()
        
    def init_ui(self):
        layout = QVBoxLayout(self)
        
        # 工具栏
        toolbar = QHBoxLayout()
        
        self.create_job_button = QPushButton("创建工作")
        self.assign_files_button = QPushButton("分配文件")
        self.close_job_button = QPushButton("关闭工作")
        self.user_manager_button = QPushButton("用户管理")
        self.model_manager_button = QPushButton("模型管理")
        
        self.create_job_button.clicked.connect(self.create_job)
        self.assign_files_button.clicked.connect(self.assign_files)
        self.close_job_button.clicked.connect(self.close_job)
        self.user_manager_button.clicked.connect(self.manage_users)
        self.model_manager_button.clicked.connect(self.manage_models)
        
        toolbar.addWidget(self.create_job_button)
        toolbar.addWidget(self.assign_files_button)
        toolbar.addWidget(self.close_job_button)
        toolbar.addWidget(self.user_manager_button)
        toolbar.addWidget(self.model_manager_button)
        
        layout.addLayout(toolbar)
        
        # 工作列表
        self.job_list = QListWidget()
        layout.addWidget(self.job_list)
        
        # 刷新工作列表
        self.refresh_jobs()
    
    def manage_users(self):
        dialog = UserManagerDialog(self.user_manager, self)
        dialog.exec()
    
    def manage_models(self):
        dialog = ModelManagerDialog(self.data_manager, self)
        dialog.exec()
    
    def create_job(self):
        """创建新工作"""
        dialog = CreateJobDialog(self.data_manager, self)
        if dialog.exec():
            job_data = dialog.get_job_data()
            self.data_manager.create_job(job_data)
            self.refresh_jobs()
            QMessageBox.information(self, "成功", "工作创建成功")
    
    def assign_files(self):
        """分配文件给用户"""
        selected_items = self.job_list.selectedItems()
        if not selected_items:
            QMessageBox.warning(self, "警告", "请先选择一个工作")
            return
            
        job = selected_items[0].data(Qt.UserRole)
        
        dialog = AssignFilesDialog(self.data_manager, job, self)
        if dialog.exec():
            user_files = dialog.get_assigned_files()
            self.data_manager.assign_files(job["id"], user_files)
            self.refresh_jobs()
            QMessageBox.information(self, "成功", "文件分配成功")
    
    def close_job(self):
        """关闭工作"""
        selected_items = self.job_list.selectedItems()
        if not selected_items:
            QMessageBox.warning(self, "警告", "请先选择一个工作")
            return
            
        job = selected_items[0].data(Qt.UserRole)
        
        reply = QMessageBox.question(
            self, "确认", f"确定要关闭工作 '{job['name']}' 吗？这将保存所有数据并从Redis中删除。",
            QMessageBox.Yes | QMessageBox.No
        )
        
        if reply == QMessageBox.Yes:
            self.data_manager.close_job(job["id"])
            self.refresh_jobs()
            QMessageBox.information(self, "成功", "工作已关闭")
    
    def refresh_jobs(self):
        """刷新工作列表"""
        self.job_list.clear()
        jobs = self.data_manager.get_jobs()
        for job in jobs:
            item = QListWidgetItem(f"{job['name']} - {job['file_type']} - {len(job['files'])}个文件")
            item.setData(Qt.UserRole, job)
            self.job_list.addItem(item)

# 统一标注检查界面类
class AnnotationInspectionWidget(QWidget):
    def __init__(self, data_manager, user_role, parent=None):
        super().__init__(parent)
        self.data_manager = data_manager
        self.user_role = user_role
        self.current_job = None
        self.current_file = None
        self.current_annotation = None
        self.init_ui()
        
    def init_ui(self):
        main_layout = QHBoxLayout(self)
        
        # 左侧面板 - 工作和文件列表
        left_panel = QVBoxLayout()
        
        # 工作列表
        job_group = QGroupBox("工作列表")
        job_layout = QVBoxLayout(job_group)
        self.job_list = QListWidget()
        self.job_list.itemClicked.connect(self.on_job_selected)
        job_layout.addWidget(self.job_list)
        left_panel.addWidget(job_group)
        
        # 文件列表
        file_group = QGroupBox("文件列表")
        file_layout = QVBoxLayout(file_group)
        self.file_list = QListWidget()
        self.file_list.itemClicked.connect(self.on_file_selected)
        file_layout.addWidget(self.file_list)
        left_panel.addWidget(file_group)
        
        main_layout.addLayout(left_panel, 1)
        
        # 右侧面板 - 标注/检查区域和控制
        right_panel = QVBoxLayout()
        
        # 工具栏 - 根据角色显示不同按钮
        toolbar = QHBoxLayout()
        
        self.prev_button = QPushButton("上一个")
        self.next_button = QPushButton("下一个")
        self.save_button = QPushButton("保存标注")
        self.submit_comment_button = QPushButton("提交评论")
        self.status_combo = QComboBox()
        self.status_combo.addItems(["未分配", "已分配", "已经浏览", "已标注", "检查完成", "检查错误"])
        
        self.prev_button.clicked.connect(self.prev_file)
        self.next_button.clicked.connect(self.next_file)
        self.save_button.clicked.connect(self.save_annotation)
        self.submit_comment_button.clicked.connect(self.submit_comment)
        
        # 根据角色显示按钮
        if self.user_role in ["管理员", "标注员"]:
            toolbar.addWidget(self.prev_button)
            toolbar.addWidget(self.next_button)
            toolbar.addWidget(self.save_button)
        
        if self.user_role in ["管理员", "检查员"]:
            toolbar.addWidget(self.submit_comment_button)
        
        toolbar.addWidget(QLabel("状态:"))
        toolbar.addWidget(self.status_combo)
        
        right_panel.addLayout(toolbar)
        
        # 标注/检查区域
        self.annotation_area = AnnotationInspectionArea(self, self.user_role)
        right_panel.addWidget(self.annotation_area, 1)
        
        # 标签和属性控制
        attrs_layout = QGridLayout()
        
        attrs_layout.addWidget(QLabel("对象标签:"), 0, 0)
        self.obj_label_combo = QComboBox()
        attrs_layout.addWidget(self.obj_label_combo, 0, 1)
        
        attrs_layout.addWidget(QLabel("OCR转写:"), 1, 0)
        self.transliterate_edit = QLineEdit()
        attrs_layout.addWidget(self.transliterate_edit, 1, 1)
        
        attrs_layout.addWidget(QLabel("评论类型:"), 2, 0)
        self.comment_type_combo = QComboBox()
        self.comment_type_combo.addItems(["正确", "无标签", "目标不完整", "标签不正确", "不标准"])
        attrs_layout.addWidget(self.comment_type_combo, 2, 1)
        
        attrs_layout.addWidget(QLabel("评论内容:"), 3, 0)
        self.comment_text_edit = QTextEdit()
        self.comment_text_edit.setPlaceholderText("输入评论内容...")
        attrs_layout.addWidget(self.comment_text_edit, 3, 1, 1, 2)
        
        # 根据角色显示控件
        if self.user_role in ["管理员", "标注员"]:
            right_panel.addLayout(attrs_layout)
        elif self.user_role == "检查员":
            right_panel.addLayout(attrs_layout)
        
        main_layout.addLayout(right_panel, 3)
        
        # 刷新工作列表
        self.refresh_jobs()
    
    def refresh_jobs(self):
        self.job_list.clear()
        jobs = self.data_manager.get_jobs()
        for job in jobs:
            item = QListWidgetItem(f"{job['name']} - {job['file_type']} - {len(job['files'])}个文件")
            item.setData(Qt.UserRole, job)
            self.job_list.addItem(item)
    
    def on_job_selected(self, item):
        self.current_job = item.data(Qt.UserRole)
        self.refresh_files()
        
        # 更新标签下拉框
        self.obj_label_combo.clear()
        self.obj_label_combo.addItems(self.current_job.get("labels", []))
    
    def refresh_files(self):
        self.file_list.clear()
        if not self.current_job:
            return
            
        # 根据角色获取不同的文件列表
        if self.user_role == "标注员":
            # 标注员获取分配给自己的文件
            files = self.data_manager.get_job_files(
                self.current_job["id"], 
                owner=self.data_manager.username
            )
        elif self.user_role == "检查员":
            # 检查员获取所有已标注的文件
            files = self.data_manager.get_job_files(
                self.current_job["id"], 
                status="已标注"
            )
        else:  # 管理员获取所有文件
            files = self.current_job.get("files", [])
        
        for file in files:
            status_color = {
                "未分配": "gray",
                "已分配": "blue",
                "已经浏览": "purple",
                "已标注": "green",
                "检查完成": "darkgreen",
                "检查错误": "red"
            }.get(file.get("status", "未分配"), "black")
            
            item = QListWidgetItem(f"{file.get('filename', '')} [{file.get('status', '未分配')}]")
            item.setData(Qt.UserRole, file)
            item.setForeground(QColor(status_color))
            self.file_list.addItem(item)
    
    def on_file_selected(self, item):
        self.current_file = item.data(Qt.UserRole)
        self.load_annotation()
    
    def load_annotation(self):
        if not self.current_job or not self.current_file:
            return
            
        self.current_annotation = self.data_manager.get_file_annotation(
            self.current_job["id"], 
            self.current_file.get("filename", "")
        )
        
        # 更新UI
        if self.current_annotation:
            self.annotation_area.set_annotation_data(self.current_annotation)
            self.status_combo.setCurrentText(self.current_annotation.get("file_attributes", {}).get("status", "未分配"))
            self.obj_label_combo.setCurrentText(self.current_annotation.get("file_attributes", {}).get("obj", ""))
            self.transliterate_edit.setText(self.current_annotation.get("file_attributes", {}).get("transliterate", ""))
            
            # 初始化评论区域
            if self.user_role in ["管理员", "检查员"]:
                comment = ""
                for region in self.current_annotation.get("regions", []):
                    comment = region.get("region_attributes", {}).get("comment", "")
                    if comment:
                        break
                self.comment_text_edit.setText(comment)
    
    def save_annotation(self):
        if not self.current_annotation or self.user_role != "标注员":
            return
            
        # 更新标注数据
        if "file_attributes" in self.current_annotation:
            self.current_annotation["file_attributes"]["status"] = self.status_combo.currentText()
            self.current_annotation["file_attributes"]["obj"] = self.obj_label_combo.currentText()
            self.current_annotation["file_attributes"]["transliterate"] = self.transliterate_edit.text()
            
            # 更新区域属性
            for region in self.current_annotation.get("regions", []):
                region["region_attributes"]["obj"] = self.obj_label_combo.currentText()
                region["region_attributes"]["transliterate"] = self.transliterate_edit.text()
        
        # 保存
        if self.current_job and "id" in self.current_job:
            self.data_manager.save_file_annotation(
                self.current_job["id"], 
                self.current_annotation
            )
            
            # 刷新文件列表
            self.refresh_files()
            QMessageBox.information(self, "成功", "标注已保存")
    
    def submit_comment(self):
        if not self.current_annotation or self.user_role != "检查员":
            return
            
        comment_type = self.comment_type_combo.currentText()
        comment_text = self.comment_text_edit.toPlainText()
        
        if not comment_text:
            QMessageBox.warning(self, "警告", "请输入评论内容")
            return
            
        # 更新区域评论
        if "regions" in self.current_annotation:
            for region in self.current_annotation["regions"]:
                region["region_attributes"]["comment"] = f"{comment_type}: {comment_text}"
            
            # 更新文件状态
            if comment_type == "正确":
                self.current_annotation["file_attributes"]["status"] = "检查完成"
            else:
                self.current_annotation["file_attributes"]["status"] = "检查错误"
        
        # 保存
        if self.current_job and "id" in self.current_job:
            self.data_manager.save_file_annotation(
                self.current_job["id"], 
                self.current_annotation
            )
            
            # 刷新文件列表
            self.refresh_files()
            QMessageBox.information(self, "成功", "评论已提交")
    
    def prev_file(self):
        if not self.current_job or not self.current_file:
            return
            
        # 获取当前文件索引
        current_idx = self.file_list.currentRow()
        if current_idx > 0:
            self.file_list.setCurrentRow(current_idx - 1)
            self.on_file_selected(self.file_list.currentItem())
    
    def next_file(self):
        if not self.current_job or not self.current_file:
            return
            
        # 获取当前文件索引
        current_idx = self.file_list.currentRow()
        if current_idx < self.file_list.count() - 1:
            self.file_list.setCurrentRow(current_idx + 1)
            self.on_file_selected(self.file_list.currentItem())

# 统一标注检查区域
class AnnotationInspectionArea(QWidget):
    def __init__(self, parent, user_role, is_editable=True):
        super().__init__(parent)
        self.parent_widget = parent
        self.user_role = user_role
        self.annotation_data = None
        self.image = None
        self.scale_factor = 1.0
        self.drag_start = None
        self.current_rect = None
        self.rects = []
        self.is_editable = is_editable if user_role == "标注员" else False
        
    def set_annotation_data(self, annotation_data):
        self.annotation_data = annotation_data
        
        # 加载图片
        file_url = annotation_data.get("file_attributes", {}).get("file_url", "")
        self.image = QImage(file_url)
        
        # 提取已有的矩形标注
        self.rects = []
        for region in annotation_data.get("regions", []):
            if region.get("shape_attributes", {}).get("name") == "rect":
                x = region.get("shape_attributes", {}).get("x", 0)
                y = region.get("shape_attributes", {}).get("y", 0)
                width = region.get("shape_attributes", {}).get("width", 0)
                height = region.get("shape_attributes", {}).get("height", 0)
                self.rects.append((x, y, width, height))
        
        self.update()
    
    def paintEvent(self, event):
        if not self.image:
            super().paintEvent(event)
            return
            
        painter = QPainter(self)
        if not painter.isActive():
            painter.begin(self)
        painter.scale(self.scale_factor, self.scale_factor)
        
        # 绘制图片
        painter.drawImage(0, 0, self.image)
        
        # 绘制已有的矩形
        for rect in self.rects:
            x, y, width, height = rect
            pen = QPen(QColor(255, 0, 0))  # 红色
            pen.setWidth(2)
            painter.setPen(pen)
            painter.drawRect(x, y, width, height)
            
            # 绘制评论文本
            for region in self.annotation_data.get("regions", []):
                if region.get("shape_attributes", {}).get("name") == "rect":
                    comment = region.get("region_attributes", {}).get("comment", "")
                    if comment:
                        painter.setPen(QColor(0, 0, 0))
                        painter.drawText(x, y - 5, comment)
        
        # 绘制当前正在绘制的矩形（仅标注员可编辑）
        if self.is_editable and self.current_rect:
            x, y, width, height = self.current_rect
            pen = QPen(QColor(0, 255, 0))  # 绿色
            pen.setWidth(2)
            painter.setPen(pen)
            painter.drawRect(x, y, width, height)
    
    def mousePressEvent(self, event):
        if (event.button() == Qt.LeftButton and self.image and 
            self.is_editable):
            self.drag_start = event.pos()
            self.current_rect = None
    
    def mouseMoveEvent(self, event):
        if (event.buttons() & Qt.LeftButton and 
            self.drag_start and self.is_editable):
            current_pos = event.pos()
            x = min(self.drag_start.x(), current_pos.x())
            y = min(self.drag_start.y(), current_pos.y())
            width = abs(current_pos.x() - self.drag_start.x())
            height = abs(current_pos.y() - self.drag_start.y())
            self.current_rect = (x, y, width, height)
            self.update()
    
    def mouseReleaseEvent(self, event):
        if (event.button() == Qt.LeftButton and 
            self.drag_start and self.is_editable):
            if self.current_rect:
                x, y, width, height = self.current_rect
                region = {
                    "shape_attributes": {
                        "name": "rect",
                        "x": x,
                        "y": y,
                        "width": width,
                        "height": height
                    },
                    "region_attributes": {
                        "obj": self.parent_widget.obj_label_combo.currentText(),
                        "transliterate": self.parent_widget.transliterate_edit.text(),
                        "comment": ""
                    }
                }
                if "regions" not in self.annotation_data:
                    self.annotation_data["regions"] = []
                self.annotation_data["regions"].append(region)
                self.rects.append(self.current_rect)
                self.current_rect = None
                self.update()
            self.drag_start = None


if __name__ == "__main__":
    app = QApplication(sys.argv)
    
    # 初始化数据管理器
    data_manager = DataManager("admin", "管理员")
    
    # 登录对话框
    login_dialog = LoginDialog(data_manager.user_manager)
    if login_dialog.exec() == QDialog.Accepted:
        user_role = login_dialog.user_role
        username = login_dialog.username_edit.text()
        data_manager = DataManager(username, user_role)
        
        main_window = QMainWindow()
        central_widget = QWidget()
        main_layout = QVBoxLayout(central_widget)
        
        if user_role == "管理员":
            admin_widget = AdminWidget(data_manager)
            main_layout.addWidget(admin_widget)
        else:
            annotation_widget = AnnotationInspectionWidget(data_manager, user_role)
            main_layout.addWidget(annotation_widget)
        
        main_window.setCentralWidget(central_widget)
        main_window.show()
    
    sys.exit(app.exec())
